import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import numpy as np
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score, precision_recall_fscore_support
import seaborn as sns
import matplotlib.pyplot as plt
from tqdm import tqdm
import time
import json
from PIL import Image, ImageEnhance
import os
from model import DogClassifier
import gc
from tabulate import tabulate
from colorama import Fore, Style, init

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 初始化 colorama
init(autoreset=True)

def load_class_names(file_path='class_names.json'):
    """加载类别名称，如果文件不存在则创建默认映射"""
    try:
        if not os.path.exists(file_path):
            # 创建默认映射
            class_names = {str(i): f"Class_{i}" for i in range(120)}
            print(f"{Fore.YELLOW}警告：{file_path} 不存在，使用默认类别名称{Style.RESET_ALL}")
            return class_names
        
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"{Fore.RED}加载类别名称时出错: {str(e)}{Style.RESET_ALL}")
        return {str(i): f"Class_{i}" for i in range(120)}

def evaluate_model(model, test_loader, criterion, device):
    """评估模型性能"""
    model.eval()
    all_preds = []
    all_labels = []
    total_loss = 0
    total_time = 0
    num_samples = 0
    
    try:
        with torch.no_grad():
            for images, labels in tqdm(test_loader, desc='评估模型'):
                start_time = time.time()
                
                images = images.to(device)
                labels = labels.to(device)
                
                outputs = model(images)
                loss = criterion(outputs, labels)
                
                _, predicted = torch.max(outputs.data, 1)
                
                inference_time = time.time() - start_time
                total_time += inference_time
                num_samples += images.size(0)
                
                all_preds.extend(predicted.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
                total_loss += loss.item()
                
                # 清理GPU内存
                del images, labels, outputs, predicted
                torch.cuda.empty_cache()
        
        avg_loss = total_loss / len(test_loader)
        avg_inference_time = total_time / num_samples
        
        return np.array(all_preds), np.array(all_labels), avg_loss, avg_inference_time
    except Exception as e:
        print(f"{Fore.RED}评估模型时出错: {str(e)}{Style.RESET_ALL}")
        raise

def plot_confusion_matrix(y_true, y_pred, class_names):
    """绘制混淆矩阵"""
    try:
        cm = confusion_matrix(y_true, y_pred)
        plt.figure(figsize=(20, 20))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                    xticklabels=class_names, yticklabels=class_names)
        plt.title('混淆矩阵')
        plt.xlabel('预测标签')
        plt.ylabel('真实标签')
        plt.xticks(rotation=90)
        plt.yticks(rotation=0)
        plt.tight_layout()
        plt.savefig('confusion_matrix.png')
        plt.close()
        print(f"{Fore.GREEN}混淆矩阵已保存到 confusion_matrix.png{Style.RESET_ALL}")
    except Exception as e:
        print(f"{Fore.RED}绘制混淆矩阵时出错: {str(e)}{Style.RESET_ALL}")

def analyze_per_class_accuracy(y_true, y_pred, class_names):
    """分析每个类别的准确率"""
    try:
        per_class_acc = {}
        for i in range(len(class_names)):
            mask = (y_true == i)
            if mask.sum() > 0:
                acc = (y_pred[mask] == i).sum() / mask.sum()
                per_class_acc[class_names[str(i)]] = acc
        
        # 绘制每个类别的准确率
        plt.figure(figsize=(15, 5))
        plt.bar(per_class_acc.keys(), per_class_acc.values())
        plt.title('每个类别的准确率')
        plt.xticks(rotation=90)
        plt.tight_layout()
        plt.savefig('per_class_accuracy.png')
        plt.close()
        print(f"{Fore.GREEN}类别准确率图表已保存到 per_class_accuracy.png{Style.RESET_ALL}")
        
        return per_class_acc
    except Exception as e:
        print(f"{Fore.RED}分析类别准确率时出错: {str(e)}{Style.RESET_ALL}")
        return {}

def test_image_robustness(model, image_path, device):
    """测试模型在不同图像条件下的鲁棒性"""
    try:
        # 基础图像预处理
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # 加载原始图像
        image = Image.open(image_path).convert('RGB')
        original_tensor = transform(image).unsqueeze(0).to(device)
        
        variations = {
            'original': original_tensor,
            'brightness+': transform(ImageEnhance.Brightness(image).enhance(1.5)).unsqueeze(0).to(device),
            'brightness-': transform(ImageEnhance.Brightness(image).enhance(0.5)).unsqueeze(0).to(device),
            'contrast+': transform(ImageEnhance.Contrast(image).enhance(1.5)).unsqueeze(0).to(device),
            'contrast-': transform(ImageEnhance.Contrast(image).enhance(0.5)).unsqueeze(0).to(device)
        }
        
        # 添加高斯噪声
        noisy_image = np.array(image)
        noise = np.random.normal(0, 25, noisy_image.shape)
        noisy_image = np.clip(noisy_image + noise, 0, 255).astype(np.uint8)
        variations['noise'] = transform(Image.fromarray(noisy_image)).unsqueeze(0).to(device)
        
        results = {}
        model.eval()
        with torch.no_grad():
            for name, img_tensor in variations.items():
                outputs = model(img_tensor)
                probs = torch.nn.functional.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probs, 1)
                results[name] = {
                    'predicted': predicted.item(),
                    'confidence': confidence.item()
                }
                
                # 清理GPU内存
                del img_tensor, outputs, probs
                torch.cuda.empty_cache()
        
        return results
    except Exception as e:
        print(f"{Fore.RED}测试图像鲁棒性时出错: {str(e)}{Style.RESET_ALL}")
        return {}

def main():
    try:
        # 设置设备
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"{Fore.CYAN}使用设备: {device}{Style.RESET_ALL}")
        
        # 加载类别名称
        class_names = load_class_names()
        
        # 数据预处理
        test_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # 创建测试数据集和加载器
        data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'dog_data', 'test')
        if not os.path.exists(data_dir):
            print(f"{Fore.RED}警告：测试数据目录 {data_dir} 不存在{Style.RESET_ALL}")
            # 尝试使用备用目录
            data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
            if not os.path.exists(data_dir):
                print(f"{Fore.RED}警告：备用测试数据目录 {data_dir} 也不存在，创建目录...{Style.RESET_ALL}")
                os.makedirs(data_dir)
                print(f"{Fore.YELLOW}请将测试图片放入 {data_dir} 目录中，目录结构应为：{Style.RESET_ALL}")
                print("data/")
                print("  ├── 类别1/")
                print("  │   ├── 图片1.jpg")
                print("  │   └── 图片2.jpg")
                print("  └── 类别2/")
                print("      ├── 图片3.jpg")
                print("      └── 图片4.jpg")
                return
        
        print(f"{Fore.GREEN}使用测试数据目录: {data_dir}{Style.RESET_ALL}")
        
        if not any(os.path.isdir(os.path.join(data_dir, d)) for d in os.listdir(data_dir)):
            print(f"{Fore.RED}警告：{data_dir}目录中没有子文件夹（类别目录）{Style.RESET_ALL}")
            print(f"{Fore.YELLOW}请按以下结构组织测试数据：{Style.RESET_ALL}")
            print("test_dir/")
            print("  ├── 类别1/")
            print("  │   ├── 图片1.jpg")
            print("  │   └── 图片2.jpg")
            print("  └── 类别2/")
            print("      ├── 图片3.jpg")
            print("      └── 图片4.jpg")
            return
        
        test_dataset = datasets.ImageFolder(data_dir, transform=test_transform)
        test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True)
        
        print(f"{Fore.GREEN}测试数据集信息:{Style.RESET_ALL}")
        print(f"测试集大小: {len(test_dataset)}")
        print(f"测试集类别数: {len(test_dataset.classes)}")
        
        # 加载模型
        num_classes = len(class_names)
        print(f"{Fore.CYAN}模型类别数: {num_classes}{Style.RESET_ALL}")
        model = DogClassifier(num_classes=num_classes).to(device)
        
        # 加载模型权重
        model_path = 'best_model.pth'
        if not os.path.exists(model_path):
            model_path = os.path.join('model', 'best_model.pth')
            if not os.path.exists(model_path):
                raise FileNotFoundError(f"模型文件不存在，已尝试路径: best_model.pth, model/best_model.pth")
        
        print(f"{Fore.GREEN}加载模型权重: {model_path}{Style.RESET_ALL}")
        state_dict = torch.load(model_path, map_location=device)
        if isinstance(state_dict, dict) and 'model_state_dict' in state_dict:
            model.load_state_dict(state_dict['model_state_dict'])
            print(f"{Fore.GREEN}成功加载模型权重（从检查点）{Style.RESET_ALL}")
            # 打印检查点信息
            if 'epoch' in state_dict and 'best_acc' in state_dict:
                print(f"{Fore.YELLOW}检查点信息: 轮次 = {state_dict['epoch']}, 最佳准确率 = {state_dict['best_acc']:.2f}%{Style.RESET_ALL}")
        else:
            model.load_state_dict(state_dict)
            print(f"{Fore.GREEN}成功加载模型权重（直接状态字典）{Style.RESET_ALL}")
        
        criterion = nn.CrossEntropyLoss()
        
        # 打印模型结构摘要
        print(f"\n{Fore.CYAN}{Style.BRIGHT}模型结构摘要:{Style.RESET_ALL}")
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"总参数量: {total_params:,}")
        print(f"可训练参数量: {trainable_params:,}")
        print(f"参数比例: {trainable_params/total_params*100:.2f}%")
        
        # 评估模型
        print(f"\n{Fore.CYAN}{Style.BRIGHT}开始评估模型...{Style.RESET_ALL}")
        predictions, labels, avg_loss, avg_inference_time = evaluate_model(
            model, test_loader, criterion, device
        )
        
        # 计算整体指标
        accuracy = accuracy_score(labels, predictions)
        precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='weighted')
        
        # 打印基础指标
        print(f"\n{Fore.GREEN}{Style.BRIGHT}模型性能指标:{Style.RESET_ALL}")
        metrics_table = [
            ["平均损失", f"{avg_loss:.4f}"],
            ["准确率", f"{accuracy*100:.2f}%"],
            ["加权精确率", f"{precision*100:.2f}%"],
            ["加权召回率", f"{recall*100:.2f}%"],
            ["加权F1分数", f"{f1*100:.2f}%"],
            ["平均推理时间", f"{avg_inference_time*1000:.2f}ms/样本"]
        ]
        print(tabulate(metrics_table, headers=["指标", "数值"], tablefmt="fancy_grid"))
        
        # 打印详细分类报告
        print(f"\n{Fore.CYAN}{Style.BRIGHT}分类报告:{Style.RESET_ALL}")
        class_names_list = [class_names[str(i)] for i in range(len(class_names))]
        print(classification_report(labels, predictions, target_names=class_names_list))
        
        # 绘制混淆矩阵
        print(f"\n{Fore.CYAN}{Style.BRIGHT}生成混淆矩阵...{Style.RESET_ALL}")
        plot_confusion_matrix(labels, predictions, class_names_list)
        
        # 分析每个类别的准确率
        print(f"\n{Fore.CYAN}{Style.BRIGHT}分析每个类别的准确率...{Style.RESET_ALL}")
        per_class_acc = analyze_per_class_accuracy(labels, predictions, class_names)
        
        # 找出表现最好和最差的类别
        best_classes = sorted(per_class_acc.items(), key=lambda x: x[1], reverse=True)[:5]
        worst_classes = sorted(per_class_acc.items(), key=lambda x: x[1])[:5]
        
        print(f"\n{Fore.GREEN}{Style.BRIGHT}表现最好的5个类别:{Style.RESET_ALL}")
        for cls, acc in best_classes:
            print(f"- {cls}: {acc*100:.2f}%")
        
        print(f"\n{Fore.RED}{Style.BRIGHT}表现最差的5个类别:{Style.RESET_ALL}")
        for cls, acc in worst_classes:
            print(f"- {cls}: {acc*100:.2f}%")
        
        # 测试模型鲁棒性
        print(f"\n{Fore.CYAN}{Style.BRIGHT}测试模型鲁棒性...{Style.RESET_ALL}")
        # 随机选择一个类别文件夹
        class_dirs = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]
        if class_dirs:
            # 随机选择一个类别
            test_class = np.random.choice(class_dirs)
            # 获取该类别下的所有图片
            test_class_path = os.path.join(data_dir, test_class)
            test_images = [f for f in os.listdir(test_class_path) if f.endswith(('.jpg', '.jpeg', '.png'))]
            if test_images:
                # 随机选择一张图片
                test_image = np.random.choice(test_images)
                test_image_path = os.path.join(test_class_path, test_image)
                print(f"{Fore.GREEN}选择测试图片: {test_image_path}{Style.RESET_ALL}")
                
                robustness_results = test_image_robustness(model, test_image_path, device)
                print(f"\n{Fore.CYAN}{Style.BRIGHT}鲁棒性测试结果:{Style.RESET_ALL}")
                
                # 创建结果表格
                robustness_table = []
                for condition, result in robustness_results.items():
                    pred_class = class_names[str(result['predicted'])]
                    robustness_table.append([
                        condition, 
                        pred_class, 
                        f"{result['confidence']*100:.2f}%"
                    ])
                
                print(tabulate(robustness_table, headers=["变换条件", "预测类别", "置信度"], tablefmt="fancy_grid"))
                
                # 检查鲁棒性
                original_pred = robustness_results['original']['predicted']
                consistent = all(result['predicted'] == original_pred for result in robustness_results.values())
                
                if consistent:
                    print(f"{Fore.GREEN}{Style.BRIGHT}✓ 模型在所有变换条件下预测一致{Style.RESET_ALL}")
                else:
                    print(f"{Fore.YELLOW}{Style.BRIGHT}⚠ 模型在不同变换条件下预测不一致{Style.RESET_ALL}")
            else:
                print(f"{Fore.RED}在类别 {test_class} 中未找到图片{Style.RESET_ALL}")
        else:
            print(f"{Fore.RED}未找到任何类别目录{Style.RESET_ALL}")
        
        # 打印总结
        print(f"\n{Fore.GREEN}{Style.BRIGHT}{'='*30} 测试总结 {'='*30}{Style.RESET_ALL}")
        print(f"{Fore.CYAN}• 模型整体准确率: {accuracy*100:.2f}%{Style.RESET_ALL}")
        print(f"{Fore.CYAN}• 加权F1分数: {f1*100:.2f}%{Style.RESET_ALL}")
        print(f"{Fore.CYAN}• 平均推理时间: {avg_inference_time*1000:.2f}ms/样本{Style.RESET_ALL}")
        print(f"{Fore.CYAN}• 表现最好的类别: {best_classes[0][0]} ({best_classes[0][1]*100:.2f}%){Style.RESET_ALL}")
        print(f"{Fore.CYAN}• 表现最差的类别: {worst_classes[0][0]} ({worst_classes[0][1]*100:.2f}%){Style.RESET_ALL}")
        print(f"{Fore.GREEN}{Style.BRIGHT}{'='*70}{Style.RESET_ALL}")
        
        # 清理资源
        gc.collect()
        torch.cuda.empty_cache()
        
    except Exception as e:
        print(f"\n{Fore.RED}{Style.BRIGHT}程序运行出错: {str(e)}{Style.RESET_ALL}")
        raise

if __name__ == '__main__':
    main()